import dgl
import torch
import numpy as np
from scipy import sparse as sp

def positional_encoding(g, pos_enc_dim):
    """
        Graph positional encoding v/ Laplacian eigenvectors
        图位置编码
    """

    # Laplacian
    # A是邻接矩阵
    A = g.adjacency_matrix_scipy(return_edge_ids=False).astype(float)
    # N是度矩阵的-1/2次方
    N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=float)
    # 求拉普拉斯矩阵 L = I - N * A * N = I - D^-1/2 * A * D^-1/2
    L = sp.eye(g.number_of_nodes()) - N * A * N

    # Eigenvectors with numpy
    # 计算矩阵的特征值和特征向量
    EigVal, EigVec = np.linalg.eig(L.toarray())
    idx = EigVal.argsort()  # increasing order 对特征值升序排序
    EigVal, EigVec = EigVal[idx], np.real(EigVec[:, idx])  # 根据排序下标进行排序, 特征向量按列排序
    g.ndata['pos_enc'] = torch.from_numpy(EigVec[:, 1:pos_enc_dim + 1]).float()  # 取前pos_enc_dim个特征值和特征向量

    # # Eigenvectors with scipy
    # EigVal, EigVec = sp.linalg.eigs(L, k=pos_enc_dim+1, which='SR')
    # EigVec = EigVec[:, EigVal.argsort()] # increasing order
    # g.ndata['pos_enc'] = torch.from_numpy(np.abs(EigVec[:,1:pos_enc_dim+1])).float()

    return g

if __name__ == '__main__':
    graph = dgl.DGLGraph()
    num_nodes = 4
    graph.add_nodes(num_nodes)
    graph.add_edges([0, 1, 1, 2], [1, 2, 3, 3])
    print(graph)
    graph = positional_encoding(graph, 2)
    print(graph)